import torch
import torch.nn as nn
from torchvision import transforms, datasets
import matplotlib.pyplot as plt
import os
import torch.optim as optim
from model import resnet34, resnet101
from Hilo_Resnet import Model as hilo_resnet
import numpy as np
from sklearn.metrics import roc_auc_score, confusion_matrix
from vgg import vgg16_bn
from ViT import vit_base_patch16_224_in21k
from Hilo_vgg import vgg16_hilo
from Fourier import hilo
import os 
import argparse
from Traditional_Fourier import hilo as traditional_hilo


if __name__ == '__main__':
    torch.manual_seed(42)
    np.random.seed(42)
    # 设置设备相关的种子（如果使用 GPU）
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(42)
    torch.autograd.set_detect_anomaly(True)

    parser = argparse.ArgumentParser()
    parser.add_argument('--modelname', type=str,  default = 'Hilo')
    parser.add_argument('--pretrain', type=bool,  default = False)
    args = parser.parse_args()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(device)

    data_transform = transforms.Compose([transforms.ToTensor()])
    data_root = os.getcwd()
    image_path = "./archive/"  # flower data set path
    if not os.path.exists('./model/'):
        os.makedirs('./model/')
    
    batch_size = 16


    train_dataset = datasets.ImageFolder(root=image_path + "train", transform=data_transform)
    train_loader = torch.utils.data.DataLoader(train_dataset,
                                            batch_size=batch_size, shuffle=True,
                                            num_workers=0)

    validate_dataset = datasets.ImageFolder(root=image_path + "val", transform=data_transform)
    validate_loader = torch.utils.data.DataLoader(validate_dataset,
                                                batch_size=batch_size, shuffle=False,
                                                num_workers=0)

    

    modelname = args.modelname
    model = {'Hilo':hilo, 'Resnet34':resnet34, 'Vgg16':vgg16_bn, 'ViT':vit_base_patch16_224_in21k, 'Hilo_vgg':vgg16_hilo, 
             'Hilo_Resnet': hilo_resnet, 'Traditional_Hilo':traditional_hilo}
   
    try:
        net = model[modelname]()
    except:
        net = model[modelname]

    if args.pretrain:
        if modelname+'.pth' in os.listdir('./model/'):
            net.load_state_dict(torch.load('./model/'+ modelname+'.pth'))
            print('读预训练模型')


    net.to(device)

    loss_function = nn.CrossEntropyLoss()
    optimizer = optim.Adam(net.parameters(), lr=0.0001)

    best_acc = 0.0

    for epoch in range(300):
        # train
        net.train()
        for step, data in enumerate(train_loader, start=0):
            
            images, labels = data
            inputs = images.to(device)
            optimizer.zero_grad()
        
            logits = net(images.to(device))
            loss = loss_function(logits, labels.to(device))
            loss.backward(retain_graph=True)
         
            optimizer.step()

            rate = (step+1)/len(train_loader)
            a = "*" * int(rate * 50)
            b = "." * int((1 - rate) * 50)
            print("\rtrain loss: {:^3.0f}%[{}->{}]{:.4f}".format(int(rate*100), a, b, loss), end="")
        print()

        # validate
        net.eval()
        acc = 0.0  # accumulate accurate number / epoch

        target_data = next(iter(validate_loader))[0].to(device)
        with torch.no_grad():
            """
            calculate loss and OA
            """
            for val_data in validate_loader:
                val_images, val_labels = val_data
                outputs = net(val_images.to(device))  # eval model only have last output layer
                # loss = loss_function(outputs, test_labels)
                predict_y = torch.max(outputs, dim=1)[1]
                acc += (predict_y == val_labels.to(device)).sum().item()
            val_accurate = acc / len(validate_dataset)
   
            if val_accurate > best_acc:
                best_acc = val_accurate
                torch.save(net.state_dict(), './model/{}.pth'.format(modelname))

            print("[epoch %d]  OA: %.4f  Best_acc: %.4f" %
                (epoch + 1, val_accurate, best_acc))

    print('Finished Training!')
